import pandas as pd
import torch

from titanic.TitanicDataset import TitanicDataset
from titanic.train import NeuralNetwork

# 读取原始的数据集
rawData = pd.read_csv("data/test.csv")
id = rawData["PassengerId"].values.astype(int)

# 定义数据集
dataset = TitanicDataset("data/test.csv")

# 加载模型
model = NeuralNetwork()
model.load_state_dict(torch.load("models/best.pth"))
# 归一化,使得每次预测的结果一致
model.eval()

# 总数据
totalData = []
idx = 0
# 获得检测的结果
for item in dataset.x_data:
    if model(item).item() > 0.5:
        totalData.append([id[idx], 1])
    else:
        totalData.append([id[idx], 0])
    idx += 1

# 记录到文件
df = pd.DataFrame(totalData, columns=["PassengerId", "Survived"])
df.to_csv("data/out.csv", index=False)
